{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T11:07:30.240499Z",
     "start_time": "2020-09-20T11:07:28.647308Z"
    },
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "fen = open('train.tags.zh-en.en', encoding='utf8')\n",
    "fzh = open('train.tags.zh-en.zh', encoding='utf8')\n",
    "en_zh = []\n",
    "while True:\n",
    "    lz = fzh.readline()\n",
    "    le = fen.readline()\n",
    "    if not lz:\n",
    "        assert not le\n",
    "        break\n",
    "    lz, le = lz.strip(), le.strip()\n",
    "    if lz.startswith('<url>'):\n",
    "        assert le.startswith('<url>')\n",
    "        lz = fzh.readline()\n",
    "        le = fen.readline()\n",
    "        assert lz.startswith('<keywords>')\n",
    "        assert le.startswith('<keywords>')\n",
    "        lz = fzh.readline()\n",
    "        le = fen.readline()\n",
    "        assert lz.startswith('<speaker>')\n",
    "        assert le.startswith('<speaker>')\n",
    "        lz = fzh.readline()\n",
    "        le = fen.readline()\n",
    "        assert lz.startswith('<talkid>')\n",
    "        assert le.startswith('<talkid>')\n",
    "        lz = fzh.readline()\n",
    "        le = fen.readline()\n",
    "        assert lz.startswith('<title>')\n",
    "        assert le.startswith('<title>')\n",
    "        lz = fzh.readline()\n",
    "        le = fen.readline()\n",
    "        assert lz.startswith('<description>')\n",
    "        assert le.startswith('<description>')\n",
    "    else:\n",
    "        if not lz:\n",
    "            assert not le\n",
    "            break\n",
    "        lee = []\n",
    "        for w in le.split(' '):\n",
    "            w = w.replace('.', '').replace(',', '')\n",
    "            if w:\n",
    "                lee.append(w)\n",
    "        en_zh.append([lee, list(lz)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T11:07:46.548430Z",
     "start_time": "2020-09-20T11:07:45.911254Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████| 48239/48239 [00:00<00:00, 77283.70it/s]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "en_words = set()\n",
    "zh_words = set()\n",
    "for s in tqdm(en_zh):\n",
    "    for w in s[0]:\n",
    "        w = w.replace('.', '').replace(',', '').lower()\n",
    "        if w:\n",
    "            en_words.add(w)\n",
    "    for w in s[1]:\n",
    "        if w:\n",
    "            zh_words.add(w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T11:07:47.729849Z",
     "start_time": "2020-09-20T11:07:47.717854Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "34734"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(en_words)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T07:27:33.415844Z",
     "start_time": "2020-09-20T07:27:33.404851Z"
    }
   },
   "outputs": [],
   "source": [
    "en_wl = ['<sos>', '<eos>', '<pad>'] + list(en_words)\n",
    "zh_wl = ['<sos>', '<eos>', '<pad>'] + list(zh_words)\n",
    "pad_id = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T07:27:33.820108Z",
     "start_time": "2020-09-20T07:27:33.788128Z"
    }
   },
   "outputs": [],
   "source": [
    "en2id = {}\n",
    "zh2id = {}\n",
    "for i, w in enumerate(en_wl):\n",
    "    en2id[w] = i\n",
    "for i, w in enumerate(zh_wl):\n",
    "    zh2id[w] = i"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T07:27:34.339423Z",
     "start_time": "2020-09-20T07:27:34.284457Z"
    }
   },
   "outputs": [],
   "source": [
    "import random\n",
    "random.shuffle(en_zh)\n",
    "dl = len(en_zh)\n",
    "train_set = en_zh[:int(dl*0.8)]\n",
    "dev_set = en_zh[int(dl*0.8):]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T07:29:40.712714Z",
     "start_time": "2020-09-20T07:29:40.689729Z"
    },
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "batch_size = 16\n",
    "data_workers = 0\n",
    "\n",
    "class MyDataSet(torch.utils.data.Dataset):\n",
    "    def __init__(self, examples):\n",
    "        self.examples = examples\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.examples)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        example = self.examples[index]\n",
    "        s1 = example[0]\n",
    "        s2 = example[1]\n",
    "        l1 = len(s1)\n",
    "        l2 = len(s2)\n",
    "        return s1, l1, s2, l2, index\n",
    "\n",
    "def the_collate_fn(batch):\n",
    "    src = [[0]*batch_size]\n",
    "    tar = [[0]*batch_size]\n",
    "    src_max_l = 0\n",
    "    for b in batch:\n",
    "        src_max_l = max(src_max_l, b[1])\n",
    "    tar_max_l = 0\n",
    "    for b in batch:\n",
    "        tar_max_l = max(tar_max_l, b[3])\n",
    "    for i in range(src_max_l):\n",
    "        l = []\n",
    "        for x in batch:\n",
    "            if i < x[1]:\n",
    "                l.append(en2id[x[0][i]])\n",
    "            else:\n",
    "                l.append(pad_id)\n",
    "        src.append(l)\n",
    "    \n",
    "    for i in range(tar_max_l):\n",
    "        l = []\n",
    "        for x in batch:\n",
    "            if i < x[3]:\n",
    "                l.append(zh2id[x[2][i]])\n",
    "            else:\n",
    "                l.append(pad_id)\n",
    "        tar.append(l)\n",
    "    indexs = [b[4] for b in batch]\n",
    "    src.append([1] * batch_size)\n",
    "    tar.append([1] * batch_size)\n",
    "    s1 = torch.LongTensor(src)\n",
    "    s2 = torch.LongTensor(tar)\n",
    "    return s1, s2, indexs\n",
    "\n",
    "train_dataset = MyDataSet(train_set)\n",
    "train_data_loader = torch.utils.data.DataLoader(\n",
    "    train_dataset,\n",
    "    batch_size=batch_size,\n",
    "    shuffle = True,\n",
    "    num_workers=data_workers,\n",
    "    collate_fn=the_collate_fn,\n",
    ")\n",
    "\n",
    "\n",
    "\n",
    "dev_dataset = MyDataSet(dev_set)\n",
    "dev_data_loader = torch.utils.data.DataLoader(\n",
    "    dev_dataset,\n",
    "    batch_size=batch_size,\n",
    "    shuffle = True,\n",
    "    num_workers=data_workers,\n",
    "    collate_fn=the_collate_fn,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T07:30:17.580226Z",
     "start_time": "2020-09-20T07:30:17.564238Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "class Encoder(nn.Module):\n",
    "    def __init__(self, input_dim, emb_dim, hid_dim, n_layers, dropout):\n",
    "        super().__init__()\n",
    "        self.hid_dim = hid_dim\n",
    "        self.n_layers = n_layers\n",
    "        self.embedding = nn.Embedding(input_dim, emb_dim)\n",
    "        self.rnn = nn.LSTM(emb_dim, hid_dim, n_layers, dropout = dropout)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, src):\n",
    "        #src = [src len, batch size]\n",
    "        embedded = self.dropout(self.embedding(src))\n",
    "        #embedded = [src len, batch size, emb dim]\n",
    "        outputs, (hidden, cell) = self.rnn(embedded)\n",
    "        #outputs = [src len, batch size, hid dim * n directions]\n",
    "        #hidden = [n layers * n directions, batch size, hid dim]\n",
    "        #cell = [n layers * n directions, batch size, hid dim]\n",
    "        #outputs are always from the top hidden layer\n",
    "        return hidden, cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T07:30:18.036314Z",
     "start_time": "2020-09-20T07:30:18.028321Z"
    }
   },
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, output_dim, emb_dim, hid_dim, n_layers, dropout):\n",
    "        super().__init__()\n",
    "        self.output_dim = output_dim\n",
    "        self.hid_dim = hid_dim\n",
    "        self.n_layers = n_layers\n",
    "        self.embedding = nn.Embedding(output_dim, emb_dim)\n",
    "        self.rnn = nn.LSTM(emb_dim, hid_dim, n_layers, dropout = dropout)\n",
    "        self.fc_out = nn.Linear(hid_dim, output_dim)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, input, hidden, cell):\n",
    "        \n",
    "        #input = [batch size]\n",
    "        #hidden = [n layers * n directions, batch size, hid dim]\n",
    "        #cell = [n layers * n directions, batch size, hid dim]\n",
    "        \n",
    "        #n directions in the decoder will both always be 1, therefore:\n",
    "        #hidden = [n layers, batch size, hid dim]\n",
    "        #context = [n layers, batch size, hid dim]\n",
    "        \n",
    "        input = input.unsqueeze(0)\n",
    "        \n",
    "        #input = [1, batch size]\n",
    "        \n",
    "        embedded = self.dropout(self.embedding(input))\n",
    "        \n",
    "        #embedded = [1, batch size, emb dim]\n",
    "                \n",
    "        output, (hidden, cell) = self.rnn(embedded, (hidden, cell))\n",
    "        \n",
    "        #output = [seq len, batch size, hid dim * n directions]\n",
    "        #hidden = [n layers * n directions, batch size, hid dim]\n",
    "        #cell = [n layers * n directions, batch size, hid dim]\n",
    "        \n",
    "        #seq len and n directions will always be 1 in the decoder, therefore:\n",
    "        #output = [1, batch size, hid dim]\n",
    "        #hidden = [n layers, batch size, hid dim]\n",
    "        #cell = [n layers, batch size, hid dim]\n",
    "        \n",
    "        prediction = self.fc_out(output.squeeze(0))\n",
    "        \n",
    "        #prediction = [batch size, output dim]\n",
    "        \n",
    "        return prediction, hidden, cell"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T07:30:18.485660Z",
     "start_time": "2020-09-20T07:30:18.474664Z"
    },
    "code_folding": [
     7
    ]
   },
   "outputs": [],
   "source": [
    "class Seq2Seq(nn.Module):\n",
    "    def __init__(self, input_word_count, output_word_count, encode_dim, decode_dim, hidden_dim, n_layers, encode_dropout, decode_dropout, device):\n",
    "        super().__init__()\n",
    "        self.encoder = Encoder(input_word_count, encode_dim, hidden_dim, n_layers, encode_dropout)\n",
    "        self.decoder = Decoder(output_word_count, decode_dim, hidden_dim, n_layers, decode_dropout)\n",
    "        self.device = device\n",
    "        \n",
    "    def forward(self, src, trg, teacher_forcing_ratio = 0.5):\n",
    "        #src = [src len, batch size]\n",
    "        #trg = [trg len, batch size]\n",
    "        #teacher_forcing_ratio is probability to use teacher forcing\n",
    "        #e.g. if teacher_forcing_ratio is 0.75 we use ground-truth inputs 75% of the time\n",
    "        batch_size = trg.shape[1]\n",
    "        trg_len = trg.shape[0]\n",
    "        trg_vocab_size = self.decoder.output_dim\n",
    "        #tensor to store decoder outputs\n",
    "        outputs = torch.zeros(trg_len, batch_size, trg_vocab_size).to(self.device)\n",
    "\n",
    "        #last hidden state of the encoder is used as the initial hidden state of the decoder\n",
    "        hidden, cell = self.encoder(src)\n",
    "\n",
    "        #first input to the decoder is the <sos> tokens\n",
    "        input = trg[0,:]\n",
    "\n",
    "        for t in range(1, trg_len):\n",
    "\n",
    "            #insert input token embedding, previous hidden and previous cell states\n",
    "            #receive output tensor (predictions) and new hidden and cell states\n",
    "            output, hidden, cell = self.decoder(input, hidden, cell)\n",
    "\n",
    "            #place predictions in a tensor holding predictions for each token\n",
    "            outputs[t] = output\n",
    "\n",
    "            #decide if we are going to use teacher forcing or not\n",
    "            teacher_force = random.random() < teacher_forcing_ratio\n",
    "\n",
    "            #get the highest predicted token from our predictions\n",
    "            top1 = output.argmax(1) \n",
    "\n",
    "            #if teacher forcing, use actual next token as next input\n",
    "            #if not, use predicted token\n",
    "            input = trg[t] if teacher_force else top1\n",
    "        \n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T07:33:29.816751Z",
     "start_time": "2020-09-20T07:33:28.826960Z"
    }
   },
   "outputs": [],
   "source": [
    "source_word_count = len(en_wl)\n",
    "target_word_count = len(zh_wl)\n",
    "encode_dim = 256\n",
    "decode_dim = 256\n",
    "hidden_dim = 512\n",
    "n_layers = 2\n",
    "encode_dropout = 0.5\n",
    "decode_dropout = 0.5\n",
    "device = torch.device('cpu')\n",
    "model = Seq2Seq(source_word_count, target_word_count, encode_dim, decode_dim, hidden_dim, n_layers, encode_dropout, decode_dropout, device).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T08:31:42.968330Z",
     "start_time": "2020-09-20T08:31:42.574573Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Seq2Seq(\n",
       "  (encoder): Encoder(\n",
       "    (embedding): Embedding(37843, 256)\n",
       "    (rnn): LSTM(256, 512, num_layers=2, dropout=0.5)\n",
       "    (dropout): Dropout(p=0.5, inplace=False)\n",
       "  )\n",
       "  (decoder): Decoder(\n",
       "    (embedding): Embedding(4015, 256)\n",
       "    (rnn): LSTM(256, 512, num_layers=2, dropout=0.5)\n",
       "    (fc_out): Linear(in_features=512, out_features=4015, bias=True)\n",
       "    (dropout): Dropout(p=0.5, inplace=False)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def init_weights(m):\n",
    "    for name, param in m.named_parameters():\n",
    "        nn.init.uniform_(param.data, -0.08, 0.08)\n",
    "        \n",
    "model.apply(init_weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T08:31:57.966998Z",
     "start_time": "2020-09-20T08:31:57.950008Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The model has 20,131,759 trainable parameters\n"
     ]
    }
   ],
   "source": [
    "def count_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "\n",
    "print(f'The model has {count_parameters(model):,} trainable parameters')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T08:32:56.137845Z",
     "start_time": "2020-09-20T08:32:56.125853Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch.optim as optim\n",
    "optimizer = optim.Adam(model.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T08:33:18.530263Z",
     "start_time": "2020-09-20T08:33:18.525266Z"
    }
   },
   "outputs": [],
   "source": [
    "criterion = nn.CrossEntropyLoss(ignore_index = pad_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T08:36:28.160130Z",
     "start_time": "2020-09-20T08:36:28.141142Z"
    },
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "def train(model, iterator, optimizer, criterion, clip):\n",
    "    model.train()\n",
    "    epoch_loss = 0\n",
    "    \n",
    "    for i, batch in enumerate(iterator):\n",
    "        \n",
    "        src = batch[0]\n",
    "        trg = batch[1]\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        output = model(src, trg)\n",
    "        \n",
    "        #trg = [trg len, batch size]\n",
    "        #output = [trg len, batch size, output dim]\n",
    "        \n",
    "        output_dim = output.shape[-1]\n",
    "        \n",
    "        output = output[1:].view(-1, output_dim)\n",
    "        trg = trg[1:].view(-1)\n",
    "        \n",
    "        #trg = [(trg len - 1) * batch size]\n",
    "        #output = [(trg len - 1) * batch size, output dim]\n",
    "        \n",
    "        loss = criterion(output, trg)\n",
    "        \n",
    "        loss.backward()\n",
    "        \n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)\n",
    "        \n",
    "        optimizer.step()\n",
    "        \n",
    "        epoch_loss += loss.item()\n",
    "        \n",
    "    return epoch_loss / len(iterator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T08:33:52.510543Z",
     "start_time": "2020-09-20T08:33:52.492554Z"
    },
    "code_folding": [
     1
    ]
   },
   "outputs": [],
   "source": [
    "def evaluate(model, iterator, criterion):\n",
    "    \n",
    "    model.eval()\n",
    "    \n",
    "    epoch_loss = 0\n",
    "    \n",
    "    with torch.no_grad():\n",
    "    \n",
    "        for i, batch in enumerate(iterator):\n",
    "\n",
    "            src = batch.src\n",
    "            trg = batch.trg\n",
    "\n",
    "            output = model(src, trg, 0) #turn off teacher forcing\n",
    "\n",
    "            #trg = [trg len, batch size]\n",
    "            #output = [trg len, batch size, output dim]\n",
    "\n",
    "            output_dim = output.shape[-1]\n",
    "            \n",
    "            output = output[1:].view(-1, output_dim)\n",
    "            trg = trg[1:].view(-1)\n",
    "\n",
    "            #trg = [(trg len - 1) * batch size]\n",
    "            #output = [(trg len - 1) * batch size, output dim]\n",
    "\n",
    "            loss = criterion(output, trg)\n",
    "            \n",
    "            epoch_loss += loss.item()\n",
    "        \n",
    "    return epoch_loss / len(iterator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T08:34:41.569431Z",
     "start_time": "2020-09-20T08:34:41.565433Z"
    }
   },
   "outputs": [],
   "source": [
    "def epoch_time(start_time, end_time):\n",
    "    elapsed_time = end_time - start_time\n",
    "    elapsed_mins = int(elapsed_time / 60)\n",
    "    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))\n",
    "    return elapsed_mins, elapsed_secs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2020-09-20T08:36:29.609Z"
    }
   },
   "outputs": [],
   "source": [
    "import math\n",
    "import time\n",
    "N_EPOCHS = 10\n",
    "CLIP = 1\n",
    "\n",
    "best_valid_loss = float('inf')\n",
    "\n",
    "for epoch in range(N_EPOCHS):\n",
    "    \n",
    "    start_time = time.time()\n",
    "    \n",
    "    train_loss = train(model, train_data_loader, optimizer, criterion, CLIP)\n",
    "    valid_loss = evaluate(model, dev_data_loader, criterion)\n",
    "    \n",
    "    end_time = time.time()\n",
    "    \n",
    "    epoch_mins, epoch_secs = epoch_time(start_time, end_time)\n",
    "    \n",
    "    if valid_loss < best_valid_loss:\n",
    "        best_valid_loss = valid_loss\n",
    "        torch.save(model.state_dict(), 'tut1-model.pt')\n",
    "    \n",
    "    print(f'Epoch: {epoch+1:02} | Time: {epoch_mins}m {epoch_secs}s')\n",
    "    print(f'\\tTrain Loss: {train_loss:.3f} | Train PPL: {math.exp(train_loss):7.3f}')\n",
    "    print(f'\\t Val. Loss: {valid_loss:.3f} |  Val. PPL: {math.exp(valid_loss):7.3f}')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
